package com.lw.leetcode.tree.b;

import com.lw.leetcode.tree.TreeNode;

import java.util.ArrayList;
import java.util.List;

/**
 * Created with IntelliJ IDEA.
 * <p>
 * tree
 * 863. 二叉树中所有距离为 K 的结点
 *
 * @author liw
 * @version 1.0
 * @date 2021/7/28 15:59
 */
public class DistanceK {

    public static void main(String[] args) {
        DistanceK test = new DistanceK();

        // [0,2,1,null,null,3]
        //3

        TreeNode instance = TreeNode.getInstance9();
        List<Integer> list = test.distanceK(instance, new TreeNode(3), 3);
        System.out.println(list);
    }


    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        List<TreeNode> list = new ArrayList<>();
        List<Integer> values = new ArrayList<>();
        find(root, target, list);
        int size = list.size() - 1;
        TreeNode node;
        addValues(list.get(size), values, 0, k);
        int count = 2;
        for (int i = size; i > 0; i--) {
            if (count - 1 == k) {
                values.add(list.get(i - 1).val);
                break;
            }
            node = list.get(i);
            TreeNode treeNode = list.get(i - 1);
            if (treeNode.left == null || treeNode.right == null) {
                count++;
                continue;
            }
            if (node.val == treeNode.left.val) {
                addValues(treeNode.right, values, count++, k);
            } else {
                addValues(treeNode.left, values, count++, k);
            }
        }
        return values;
    }

    private void addValues(TreeNode node, List<Integer> values, int n, int k) {
        if (node == null) {
            return;
        }
        if (n == k) {
            values.add(node.val);
            return;
        }
        addValues(node.left, values, n + 1, k);
        addValues(node.right, values, n + 1, k);
    }


    private boolean find(TreeNode node, TreeNode target, List<TreeNode> list) {
        if (node == null) {
            return false;
        }
        list.add(node);
        if (target.val == node.val) {
            return true;
        }
        boolean b = find(node.left, target, list);
        if (b) {
            return b;
        }
        b = find(node.right, target, list);
        if (!b) {
            list.remove(list.size() - 1);
        }
        return b;
    }
}
